import sqlite3
import csv
import ipaddress
import sys
import os

# 定义数据库名称和文件路径
db_name = '/opt/shells/dbip/dbip.db'
file_path = '/opt/shells/dbip/csv/dbip-city-lite-2024-12.csv'
if len(sys.argv)>1: 
    file_path =os.path.join('/opt/shells/dbip/csv',sys.argv[1])
print(f"{file_path}")

# 连接到 SQLite 数据库（如果数据库不存在，则会自动创建）
conn = sqlite3.connect(db_name)
cursor = conn.cursor()
cursor.execute('DROP TABLE IF EXISTS dbip_lookup;')

# 创建表（如果表不存在）
cursor.execute('''
CREATE TABLE IF NOT EXISTS dbip_lookup (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    addr_type TEXT NOT NULL,
    ip_start TEXT NOT NULL,
    ip_end TEXT NOT NULL,
    continent TEXT NOT NULL,
    country TEXT NOT NULL,
    stateprov TEXT NOT NULL,
    city TEXT NOT NULL,
    latitude REAL NOT NULL,
    longitude REAL NOT NULL
);
''')

# 准备插入的 SQL 语句
insert_query = '''
INSERT INTO dbip_lookup (addr_type, ip_start, ip_end, continent, country, stateprov, city, latitude, longitude)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?);
'''

# 打开文件并读取内容
with open(file_path, mode='r', encoding='utf-8') as file:
    reader = csv.reader(file)
    
    # 跳过CSV文件的标题行
    # next(reader, None)

    # 过滤并插入只包含 IPv4 的数据
    data_to_insert = []
    for row in reader:
        try:
            # 检查国家是否为中国
            if row[3] != 'CN':
                continue
            ip_start = ipaddress.ip_address(row[0])
            ip_end = ipaddress.ip_address(row[1])
            
            # 判断 IP 地址类型并设置 addr_type
            if ip_start.version == 4 and ip_end.version == 4:
                addr_type = 'IPv4'
            elif ip_start.version == 6 and ip_end.version == 6:
                addr_type = 'IPv6'
            else:
                continue  # 如果一个IP是v4另一个是v6则跳过该行

            # 构建要插入的数据行
            data_row = [addr_type] + row[:2] + row[2:]
            data_to_insert.append(data_row)
        except ValueError as e:
            print(f"Invalid IP address in row: {row}. Error: {e}")
            continue

    # 使用 executemany 方法批量插入数据
    cursor.executemany(insert_query, data_to_insert)

# 提交更改并关闭连接
conn.commit()
conn.close()

print(f"数据已成功插入到 {db_name} 中")